package com.study.consumer.filter;

import com.study.consumer.Constant;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.common.threadlocal.InternalThreadLocal;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.RpcContext;
import org.apache.dubbo.rpc.cluster.interceptor.ClusterInterceptor;
import org.apache.dubbo.rpc.cluster.support.AbstractClusterInvoker;

/**
 * @Auther: luobiao
 * @Date: 2021/4/29 06:26
 * @Description: 用来处理每次调用rpc服务后RpcContext#LOCAL被清除的导致tag丢失问题
 */
@Activate
public class TagInterceptor implements ClusterInterceptor {
    private static final InternalThreadLocal<String> TAG = new InternalThreadLocal<String>() {
        @Override
        protected String initialValue() {
            return "";
        }
    };

    @Override
    public void before(AbstractClusterInvoker<?> clusterInvoker, Invocation invocation) {
        String tag = RpcContext.getContext().getAttachment(Constant.DUBBO_TAG);
        TAG.set(tag);
    }

    @Override
    public void after(AbstractClusterInvoker<?> clusterInvoker, Invocation invocation) {
        RpcContext.getContext().setAttachment(Constant.DUBBO_TAG,TAG.get());
    }
}
